#!/usr/bin/env python3
import os
import cv2
import numpy as np
from dataclasses import dataclass
from typing import Optional, Tuple
import glob
import time  # 新增导入 time 模块
from shared_memory_reader import SharedMemoryReader
from logger import setup_logger

logger = setup_logger()

def get_depth_image(image_path):
    # 提取文件名
    img_name = os.path.basename(image_path)

    # 获取 Color 上层目录，然后构建对应的 Depth 路径
    base_dir = os.path.dirname(os.path.dirname(image_path))
    
    # 替换路径中的 Color -> Depth 或 Color1 -> Depth1
    sub_dir = os.path.basename(os.path.dirname(image_path))  # 比如 Color 或 Color1
    if "Color" in sub_dir:
        depth_sub_dir = sub_dir.replace("Color", "Depth")
    else:
        print(f"未识别的子目录: {sub_dir}")
        return None

    depth_path = os.path.join(base_dir, depth_sub_dir, img_name)

    # 检查文件是否存在
    if not os.path.exists(depth_path):
        # 尝试将 .jpg 替换为 .png
        if img_name.endswith('.jpg'):
            png_name = img_name.replace('.jpg', '.png')
            depth_path = os.path.join(base_dir, depth_sub_dir, png_name)
            if not os.path.exists(depth_path):
                print(f"深度图路径不存在: {depth_path}")
                return None

    try:
        # 使用 PIL 打开图像
        pil_img = Image.open(depth_path)

        # 处理不同格式的深度图
        if pil_img.mode == 'I;16':  # 16位小端深度图
            depth_array = np.array(pil_img, dtype=np.uint16)
            # 可选：检测是否需要字节序转换（一般不需要）
            if np.max(depth_array) > 65535:
                depth_array = depth_array.byteswap().newbyteorder()
        elif pil_img.mode == 'I':  # 32位整数深度图
            depth_array = np.array(pil_img, dtype=np.uint32)
            depth_array = (depth_array >> 16).astype(np.uint16)  # 提取高16位
        else:
            depth_array = np.array(pil_img)
            print(f"警告: 未知模式 {pil_img.mode}")

        # 返回单通道图
        if depth_array.ndim == 2:
            return depth_array
        else:
            print(f"警告: 意外通道数 {depth_array.shape}")
            return depth_array[:, :, 0]

    except Exception as e:
        print(f"深度图处理失败: {str(e)}")
        return None
    
@dataclass
class ImageData:
    """图像数据结构"""
    timestamp: str
    readable_timestamp: str
    color_img: np.ndarray
    depth_img: np.ndarray
    color_path: Optional[str] = None

class ImageProvider:
    """统一的图像提供接口"""
    
    def __init__(self, from_folder: bool = True, color_folder: str = "./data/Color1", depth_folder: str = "./data/Depth1"):
        self.from_folder = from_folder
        self.color_folder = color_folder
        self.depth_folder = depth_folder
        self.shm_reader = None
        self.color_files = []
        self.current_index = 0
        
        if from_folder:
            self.color_files = sorted(glob.glob(os.path.join(color_folder, "*.png")))
            logger.info(f"从文件夹读取，共找到 {len(self.color_files)} 个图片文件")
        else:
            self.shm_reader = SharedMemoryReader()
            logger.info("从共享内存读取")
    
    def get_next_image(self) -> Optional[ImageData]:
        """获取下一张图像"""
        if self.from_folder:
            return self._get_from_folder()
        else:
            return self._get_from_memory()
    
    def _get_from_folder(self) -> Optional[ImageData]:
        """从文件夹获取图像，按文件名中的数字大小排序"""
        # 按文件名中的数字排序
        self.color_files.sort(key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))
        
        if self.current_index >= len(self.color_files):
            return None
        
        color_path = self.color_files[self.current_index]
        depth_path = os.path.join(self.depth_folder, os.path.basename(color_path))
        self.current_index += 1
        
        if not os.path.exists(depth_path):
            logger.warning(f"缺失深度图: {depth_path}")
            return self._get_from_folder()  # 递归获取下一张
        
        color_img = cv2.imread(color_path)
        depth_img = get_depth_image(color_path)
        
        if color_img is None or depth_img is None:
            logger.warning(f"读取图像失败: {color_path} 或 {depth_path}")
            return self._get_from_folder()  # 递归获取下一张
        
        timestamp = os.path.splitext(os.path.basename(color_path))[0]
        
        return ImageData(
            timestamp=timestamp,
            readable_timestamp=timestamp,
            color_img=color_img,
            depth_img=depth_img,
            color_path=color_path
        )

    def _get_from_memory(self) -> Optional[ImageData]:
        """从共享内存获取图像"""
        logger.info("开始从共享内存读取图像")
        start_time = time.time()  # 记录开始时间
        
        readable_timestamp, timestamp, color_img, depth_img = self.shm_reader.read_images()
        
        read_time = time.time() - start_time  # 计算读取耗时
        logger.info(f"从共享内存读取图像完成，耗时: {read_time:.4f} 秒")
        
        if color_img is None or depth_img is None:
            logger.warning("从共享内存读取图像失败")
            return None
        
        # 可选：保存图像到文件
        write_start_time = time.time()  # 记录写入开始时间
        success, color_path, depth_path = self.shm_reader.save_images(timestamp, color_img, depth_img)
        write_time = time.time() - write_start_time  # 计算写入耗时
        
        if success:
            logger.info(f"图像保存成功，路径: {color_path}, {depth_path}, 写入耗时: {write_time:.4f} 秒")
        else:
            logger.warning("图像保存失败")
        
        return ImageData(
            timestamp=timestamp,
            readable_timestamp=readable_timestamp,
            color_img=color_img,
            depth_img=depth_img,
            color_path=color_path if success else None
        )
    
    def close(self):
        """关闭资源"""
        if self.shm_reader:
            self.shm_reader.close()